# SENet 改进版

import torch
from torch import nn
import math


class ECANet(nn.Module):
    def __init__(self, channels, ratio=16, gamma=2, b=1):
        super(ECANet, self).__init__()
        # 输入通道数自适应修改卷积核大小
        kernel_size = int(abs((math.log(channels, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
        padding = kernel_size // 2
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size, padding=padding, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()
        # b c h w     b c 1 1
        avg = self.avg_pool(x).view([b, 1, c])
        conv_out = self.conv(avg)
        out = self.sigmoid(conv_out).view([b, c, 1, 1])
        return out * x
